"""
    @file Sklearn
    @description Decision Tree By Sklearn Using Entropy
    @author Synhard
    @tel 13001321080
    @id 21126338
    @email 823436512@qq.com
    @date 2021-09-24 18:03
    @version 1.0
"""
import matplotlib
import numpy as np
import pandas as pd
import random

from matplotlib import pyplot as plt
from sklearn.metrics import confusion_matrix

matplotlib.rcParams['font.family'] = 'SimHei'  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")
random.seed(0)


def train_test_split(df, test_size):
    if isinstance(test_size, float):
        test_size = round(test_size * len(df))

    indices = df.index.tolist()
    test_indices = random.sample(population=indices, k=test_size)

    test_df = df.loc[test_indices]
    train_df = df.drop(test_indices)

    return train_df, test_df


def check_purity(data):
    label_column = data[:, -1]
    unique_classes = np.unique(label_column)

    if len(unique_classes) == 1:
        return True
    else:
        return False


def classify_data(data):
    label_column = data[:, -1]
    unique_classes, counts_unique_classes = np.unique(label_column, return_counts=True)

    index = counts_unique_classes.argmax()
    classification = unique_classes[index]

    return classification


def get_potential_splits(data):
    potential_splits = {}
    _, n_columns = data.shape
    for column_index in range(n_columns - 1):
        potential_splits[column_index] = []
        values = data[:, column_index]
        unique_values = np.unique(values)

        for index in range(len(unique_values)):
            if index != 0:
                current_value = unique_values[index]
                previous_value = unique_values[index - 1]
                potential_split = (current_value + previous_value) / 2

                potential_splits[column_index].append(potential_split)

    return potential_splits


def split_data(data, split_column, split_value):
    split_column_values = data[:, split_column]

    data_below = data[split_column_values <= split_value]
    data_above = data[split_column_values > split_value]

    return data_below, data_above


def calculate_entropy(data):
    label_column = data[:, -1]
    _, counts = np.unique(label_column, return_counts=True)

    probabilities = counts / counts.sum()
    entropy = sum(probabilities * -np.log2(probabilities))

    return entropy


def calculate_overall_entropy(data_below, data_above):
    n = len(data_below) + len(data_above)
    p_data_below = len(data_below) / n
    p_data_above = len(data_above) / n

    overall_entropy = (p_data_below * calculate_entropy(data_below)
                       + p_data_above * calculate_entropy(data_above))

    return overall_entropy


def determine_best_split(data, potential_splits):
    overall_entropy = 1000
    best_split_column = 0
    best_split_value = 0
    for column_index in potential_splits:
        for value in potential_splits[column_index]:
            data_below, data_above = split_data(data, split_column=column_index, split_value=value)
            current_overall_entropy = calculate_overall_entropy(data_below, data_above)

            if current_overall_entropy <= overall_entropy:
                overall_entropy = current_overall_entropy
                best_split_column = column_index
                best_split_value = value

    return best_split_column, best_split_value


def decision_tree_algorithm(df, counter=0, min_samples=2, max_depth=4):
    if counter == 0:
        global COLUMN_HEADERS
        COLUMN_HEADERS = df.columns
        data = df.values
    else:
        data = df
    if (check_purity(data)) or (len(data) < min_samples) or (counter == max_depth):
        classification = classify_data(data)

        return classification
    else:
        counter += 1

        potential_splits = get_potential_splits(data)
        split_column, split_value = determine_best_split(data, potential_splits)

        data_below, data_above = split_data(data, split_column, split_value)
        feature_name = COLUMN_HEADERS[split_column]
        question = "{} <= {}".format(feature_name, split_value)
        sub_tree = {question: []}
        yes_answer = decision_tree_algorithm(data_below, counter, min_samples, max_depth)
        no_answer = decision_tree_algorithm(data_above, counter, min_samples, max_depth)
        if yes_answer == no_answer:
            sub_tree = yes_answer
        else:
            sub_tree[question].append(yes_answer)
            sub_tree[question].append(no_answer)

        return sub_tree


def classify_example(example, tree):
    question = list(tree.keys())[0]
    feature_name, comparison_operator, value = question.split(" ")
    if example[feature_name] <= float(value):
        answer = tree[question][0]
    else:
        answer = tree[question][1]

    if not isinstance(answer, dict):
        return answer
    else:
        residual_tree = answer
        return classify_example(example, residual_tree)


def calculate_accuracy(df, tree):
    df["classification"] = df.apply(classify_example, args=(tree,), axis=1)
    df["classification_correct"] = df["classification"] == df["target"]
    accuracy = df["classification_correct"].mean()

    return accuracy


def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for obj in secondDict:
        if type(obj).__name__ == 'dict':
            # nodes
            numLeafs += getNumLeafs(obj)
        else:
            numLeafs += 1
    return numLeafs


def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for obj in secondDict:
        if type(obj).__name__ == 'dict':
            thisDepth = 1 + getTreeDepth(obj)
        else:
            thisDepth = 1
        if thisDepth > maxDepth: maxDepth = thisDepth
    return maxDepth


def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                            xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)


def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    firstStr = list(myTree.keys())[0]
    cntrPt = (plotTree.xOff + (2 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yOff)
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD
    for i in range(0, len(secondDict)):
        if type(secondDict[i]).__name__ == 'dict':
            if i == 0:
                plotTree(secondDict[i], cntrPt, "Y")  # 递归
            else:
                plotTree(secondDict[i], cntrPt, "N")  # 递归
        else:
            plotTree.xOff = plotTree.xOff + 1.1 / plotTree.totalW
            plotNode(secondDict[i], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
            if i == 0:
                plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, "Y")  # 递归
            else:
                plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, "N")  # 递归

    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD


def createPlot(myTree):
    fig = plt.figure(1, facecolor='white', figsize=(8, 8))
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotTree.totalW = float(getNumLeafs(myTree))
    plotTree.totalD = float(getTreeDepth(myTree))
    plotTree.xOff = -1.5 / plotTree.totalW
    plotTree.yOff = 1.1
    plotTree(myTree, (0.5, 1.0), '')
    plt.show()


def prediction(decisontree, row):
    if isinstance(decisontree, float):
        return decisontree
    root = list(decisontree.keys())[0]
    attr = root.split(" <= ")[0]
    value = root.split(" <= ")[1]
    if float(row.get(attr)) <= float(value):
        return prediction(decisontree.get(root)[0], row)
    else:
        return prediction(decisontree.get(root)[1], row)
    pass


def testTree(decisontree, dataSet):
    testmemory = []
    label = []
    TP = 0
    FP = 0
    TN = 0
    FN = 0
    rowDict = []
    for i in range(1, dataSet.shape[0]):
        dict = {}
        for j in range(0, dataSet[0].shape[0]):
            dict[dataSet[0][j]] = dataSet[i][j]

        rowDict.append(dict)

    for row in rowDict:
        label.append(row['target'])
        testmemory.append(str(prediction(decisontree, row)))
        pass

    for i in range(0, len(rowDict)):
        if (testmemory[i] == '1.0') and (label[i] == '1.0'):  # testmemory为预测 label为实际
            TP += 1
        elif (testmemory[i] == '1.0') and (label[i] == '0.0'):
            FP += 1
        elif (testmemory[i] == '0.0') and (label[i] == '0.0'):
            TN += 1
        elif (testmemory[i] == '0.0') and (label[i] == '1.0'):
            FN += 1

    print('TP:%d' % TP)  # 真阳性
    print('FP:%d' % FP)  # 假阳性
    print('TN:%d' % TN)  # 真阴性
    print('FN:%d' % FN)  # 假阴性

    cm = confusion_matrix(label, testmemory, labels=["0.0", "1.0"])
    plt.rc('figure', figsize=(8, 8))
    plt.matshow(cm, cmap=plt.cm.cool)  # 背景颜色
    plt.colorbar()  # 颜色标签
    # 内部添加图例标签
    for x in range(len(cm)):
        for y in range(len(cm)):
            plt.annotate(cm[x, y], xy=(y, x), horizontalalignment='center', verticalalignment='center')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.title('decision_tree')
    plt.savefig(r'confusion_matrix')
    plt.show()


# 主函数
if __name__ == '__main__':
    df = pd.read_excel("heart.xlsx")
    train_df, test_df = train_test_split(df, test_size=0.3)
    data = train_df.values
    decisionTree = decision_tree_algorithm(train_df, max_depth=5)
    test_data = list(test_df.values)
    test_data.insert(0, list(test_df.columns))
    testTree(decisionTree, np.array(test_data))
    createPlot(decisionTree)
